#include <EXTERN.h>
#include <perl.h>
#include <stdio.h>
#include <stdint.h>
#include <sys/un.h>
#include <sys/wait.h>
#include <sys/types.h>
#include <grp.h>
#include <pwd.h>
#include <getopt.h>

#include "milter.h"

#ifndef DEFAULT_LISTEN
#define DEFAULT_LISTEN "unix:/var/run/milter-perl/sock"
#endif

#ifndef DEFAULT_THREADS
#define DEFAULT_THREADS 0
#endif

#ifdef DEBUG
# define Trace(...) if (debug) warn(__VA_ARGS__)
#else
# define Trace(...)
#endif

EXTERN_C void xs_init(pTHX);

static PerlInterpreter *my_perl;

static struct callback_table_t {
    uint8_t code;
    uint8_t flag;
    CV **callback;
    int callback_count;
    char *name;
    char *fmt;
} callback_table[] = {
    { SMFIC_CONNECT, SMFIP_NOCONNECT, NULL, 0, "connect", "scns"  },
    { SMFIC_HELO,    SMFIP_NOHELO,    NULL, 0, "helo",    "s"     },
    { SMFIC_MAIL,    SMFIP_NOMAIL,    NULL, 0, "envfrom", "s@"    },
    { SMFIC_RCPT,    SMFIP_NORCPT,    NULL, 0, "envrcpt", "s@"    },
    { SMFIC_HEADER,  SMFIP_NOHDRS,    NULL, 0, "header",  "ss"    },
    { SMFIC_EOH,     SMFIP_NOEOH,     NULL, 0, "eoh",     NULL    },
    { SMFIC_BODY,    SMFIP_NOBODY,    NULL, 0, "body",    "b"     },
    { SMFIC_BODYEOB, SMFIP_NOBODY,    NULL, 0, "eom",     NULL    },
    { SMFIC_ABORT,   0,               NULL, 0, "abort",   NULL    },
    { SMFIC_QUIT,    0,               NULL, 0, "close",   NULL    },
    { SMFIC_MACRO,   0,               NULL, 0, "macro",   "cs@"   },
    { 0,             0,               NULL, 0, NULL,      NULL    },
};
static struct callback_table_t *callbacks[256];
static uint32_t actions = SMFI_CURR_ACTS;
static uint32_t protocols = SMFIP_NONE;
static uint8_t error_result = SMFIR_ACCEPT;
static int hup = 0;
static int debug = 0;

static int namespace_count = 0;
static char **namespaces;

void handle_hup(__attribute__((__unused__)) int signal)
{
    hup = 1;
}

void callback_init(char *namespace)
{
    struct callback_table_t *cb;
    int i;

    for (i = 0; (cb = &callback_table[i]) && cb->name; i++)
    {
        CV *cv;

        if (!(cv = get_cv(form("%s::%s", namespace, cb->name), FALSE)))
            continue;

        Trace("found callback: %s::%s = 0x%p\n", namespace, cb->name, cv);

        if (!(cb->callback = realloc(cb->callback, (cb->callback_count + 1) * sizeof(CV *))))
            croak("realloc failed: %s", strerror(errno));

        cb->callback[cb->callback_count++] = cv;

        callbacks[cb->code] = cb;

        if (cb->flag)
            protocols &= ~cb->flag;
    }
}

/* let perl handle fork, so it keeps it's internal in order */
static int perl_fork()
{
    int ret;
    SV *sv;

    dSP;
    ENTER;
    SAVETMPS;

    PUSHMARK(SP);
    sv = newSVpv("fork", sizeof("fork")-1);
    PUTBACK;

    eval_sv(sv, G_NOARGS);

    SPAGAIN;

    ret = POPi;
    SvREFCNT_dec(sv);

    PUTBACK;
    FREETMPS;
    LEAVE;

    return ret;
}

static void callback_noargs(CV *cb)
{
    dSP;
    ENTER;
    SAVETMPS;

    PUSHMARK(SP);

    call_sv((SV *)cb, G_VOID|G_NOARGS|G_EVAL|G_DISCARD);

    SPAGAIN;

    if (SvTRUE(ERRSV))
    {
        STRLEN n_a;
        warn("callback died: %s", SvPV(ERRSV, n_a));
    }

    PUTBACK;
    FREETMPS;
    LEAVE;
}

static int callback(CV *cb, SV *state, char *fmt, char *buf, int len)
{
    int count;
    int ret = error_result;

    dSP;
    ENTER;
    SAVETMPS;

    PUSHMARK(SP);
    XPUSHs(state);
    while (fmt && *fmt)
    {
        Trace("fmt=0x%p, len=%d\n", fmt, len);

        switch (*fmt++) {
        case 'b':
            Trace("  whole buf\n");
            mXPUSHp(buf, len);
            mXPUSHi(len);
            buf += len;
            len = 0;
            break;
        case 'c':
            if (len < 1)
                croak("not enough data");
            Trace("  character=%c\n", *(char *)buf);
            mXPUSHp(buf, 1);
            buf++;
            len--;
            break;
        case 's': {
            while (len > 0)
            {
                int l = strlen(buf);
                Trace("  string(%d)=%s (len=%d)\n", l, buf, len);
                mXPUSHp(buf, l);
                buf += (l + 1);
                len -= (l + 1);
                if (*fmt != '@')
                    break;
            }
            if (*fmt == '@')
                fmt++;
            break;
        }
        case 'n':
            if (len < 2)
                croak("not enough data");
            Trace("  network short=%d\n", ntohs(*(uint16_t *)buf));
            mXPUSHi(ntohs(*(uint16_t *)buf));
            buf += 2;
            len -= 2;
            break;
        default:
            croak("len=%d, fmt=%c\n", len, *fmt);
        }
    }
    PUTBACK;

    if ((fmt && *fmt) || len > 0)
        croak("error parsing structure: fmt=0x%p, len=%d, *fmt=%c\n", fmt, len, *fmt);

    Trace("calling 0x%p\n", cb);

    count = call_sv((SV *)cb, G_SCALAR|G_EVAL);

    SPAGAIN;

    /* Check the eval first */
    if (SvTRUE(ERRSV))
    {
        STRLEN n_a;
        warn("callback died: %s", SvPV(ERRSV, n_a));

        (void)POPs;

        ret = error_result;
    }
    else if (count != 1)
    {
        warn("bad number of return values: %d", count);

        ret = error_result;
    }
    else
    {
        ret = POPi;
    }

    PUTBACK;
    FREETMPS;
    LEAVE;

    return ret;
}

int write_packet(int sock, char code, char *extra, int extralen)
{
    char buf[4+1+extralen];
    char *ptr = buf;
    int len = sizeof(buf);
    int l;

    Trace("sending code=%c extra=%s\n", code, extra);

    *(uint32_t *)buf = htonl(1 + extralen);
    buf[4] = code;
    memcpy(&buf[5], extra, extralen);

    while (len > 0)
    {
        if ((l = write(sock, ptr, len)) < 0)
            return -1;

        ptr += l;
        len -= l;
    }

    return 0;
}

int fill(int sock, char *buf, int len)
{
    int l;

    while (len > 0)
    {
        if ((l = read(sock, buf, len)) <= 0)
            return -1;

        buf += l;
        len -= l;
    }

    return 0;
}

int process_packet(SV *state)
{
    struct milter_ctx *ctx = MILTER_CTX(state);
    static char buf[1 + MILTER_CHUNK_SIZE + 1];
    uint32_t len;
    struct callback_table_t *cb;
    uint8_t result = SMFIR_CONTINUE;

    Trace("ctx = 0x%p = { sock=%d, reply=0x%p, priv=0x%p, macro=0x%p }\n",
        ctx, ctx->sock, ctx->reply, ctx->priv, ctx->macro);

    if (fill(ctx->sock, (char *)&len, sizeof(len)) < 0)
        return -2;

    len = ntohl(len);

    if (len > 1 + MILTER_CHUNK_SIZE)
        return -3;

    if (fill(ctx->sock, &buf[0], len) < 0)
        return -2;

    /* force a trailing NULL; safety precaution for fmt code */
    buf[len] = 0;

    switch (ctx->code = *(uint8_t *)buf) {
    case SMFIC_MACRO: {
        char *ptr = buf + 2;
        char *end = &buf[len];
        SV *sv[2];
        int i = 0;
        while (ptr < end)
        {
            int l = strlen(ptr);
            sv[i++] = newSVpv(ptr, l);
            ptr += l + 1;

            if (i == 2)
            {
                Trace("MACRO: %s=%s\n", SvPV_nolen(sv[0]), SvPV_nolen(sv[1]));
                hv_store_ent(ctx->macro, sv[0], SvREFCNT_inc(sv[1]), FALSE);
                SvREFCNT_dec(sv[0]);
                SvREFCNT_dec(sv[1]);
                i = 0;
            }
        }
        assert(i == 0);
        result = 0; /* suppress reply */
        break;
    }
    case SMFIC_OPTNEG:
        if (len != 13)
            croak("odd packet length for SMFIC_OPTNEG: %d", len);
        if (ntohl(*(uint32_t *)&buf[1]) != 2)
            croak("version wrong");
        *(uint32_t *)&buf[5] = htonl(ctx->actions = ntohl(*(uint32_t *)&buf[5]) & actions);
        *(uint32_t *)&buf[9] = htonl(ctx->protocols = ntohl(*(uint32_t *)&buf[9]) & protocols);
        write_packet(ctx->sock, SMFIC_OPTNEG, &buf[1], len - 1);
#define ACTION(flag) ((ctx->actions & SMFIF_##flag) ? " " #flag : "")
        Trace("OPTNEG: actions=%s%s%s%s%s%s%s\n",
            ACTION(ADDHDRS), ACTION(CHGBODY), ACTION(ADDRCPT), ACTION(DELRCPT),
            ACTION(CHGHDRS), ACTION(QUARANTINE), ACTION(SETSENDER));
#define PROTOCOL(flag) ((~ctx->protocols & SMFIP_NO##flag) ? " " #flag : "")
        Trace("OPTNEG: protocols=%s%s%s%s%s%s%s\n",
            PROTOCOL(CONNECT), PROTOCOL(HELO), PROTOCOL(MAIL), PROTOCOL(RCPT),
            PROTOCOL(BODY), PROTOCOL(HDRS), PROTOCOL(EOH));
        return 0;
        break;
    case SMFIC_QUIT:
        Trace("QUIT\n");
        return -1;
        break;
    }

    if ((cb = callbacks[ctx->code]) != NULL)
    {
        int ret;
        int i = 0;

        do {
            Trace("callback %d/%d = 0x%p\n", i + 1, cb->callback_count, cb->callback[i]);

            ret = callback(cb->callback[i], state, cb->fmt, &buf[1], len - 1);

            Trace("callback returned %d\n", ret);

            switch (ret) {
            case SMFIS_CONTINUE:
                /* don't override real result with continue */
                break;
            case SMFIS_REJECT:
                result = SMFIR_REJECT;
                break;
            case SMFIS_DISCARD:
                result = SMFIR_DISCARD;
                break;
            case SMFIS_ACCEPT:
                result = SMFIR_ACCEPT;
                break;
            case SMFIS_TEMPFAIL:
                result = SMFIR_TEMPFAIL;
                break;
            }
        } while (++i < cb->callback_count);
    }

    if (((result == SMFIR_TEMPFAIL) || (result == SMFIR_REJECT)) && ctx->reply)
        write_packet(ctx->sock, SMFIR_REPLYCODE, ctx->reply, strlen(ctx->reply)+1);
    else if (result)
        write_packet(ctx->sock, result, NULL, 0);

    if (ctx->reply)
    {
        Trace("freeing reply\n");
        free(ctx->reply);
        ctx->reply = NULL;
    }

    Trace("complete\n");

    return 0;
}

int parse_address(const char *str, struct in_addr *addr, uint16_t *port)
{
    char *end;
    struct hostent *entry;

    assert(str);
    assert(addr);
    assert(port);

    *port = (uint16_t)strtol(str, &end, 10);

    if (*end++ != '@')
        croak("couldn't parse inet:<PORT>@<ADDRESS>");

    if (!(entry = gethostbyname(end)))
        croak("couldn't resolve %s", end);

    if (entry->h_addrtype != AF_INET)
        croak("non-inet return");

    addr->s_addr = *(uint32_t *)entry->h_addr_list[0];

    return 1;
}

int bind_server_unix(char *path)
{
    struct sockaddr_un sa;
    socklen_t salen;
    int listen_sock;

    sa.sun_family = AF_UNIX;
    strcpy(sa.sun_path, path);
    salen = SUN_LEN(&sa);

    unlink(sa.sun_path);

    if ((listen_sock = socket(AF_UNIX, SOCK_STREAM, 0)) < 0)
        croak("socket: %s", strerror(errno));

    if (bind(listen_sock, (struct sockaddr *)&sa, salen) < 0)
        croak("bind: %s", strerror(errno));

    if (listen(listen_sock, SOMAXCONN) < 0)
        croak("listen: %s", strerror(errno));

    return listen_sock;
}

int bind_server_inet(struct in_addr *addr, int port)
{
    struct sockaddr_in sa;
    socklen_t salen;
    int listen_sock;

    sa.sin_family = AF_INET;
    sa.sin_port = htons(port);
    sa.sin_addr.s_addr = addr->s_addr;
    salen = sizeof(struct sockaddr_in);

    if ((listen_sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)) < 0)
        croak("socket: %s", strerror(errno));

    if (setsockopt(listen_sock, SOL_SOCKET, SO_REUSEADDR, (int[1]){ 1 }, sizeof(int)) < 0)
        warn("setsockopt(SO_REUSEADDR): %s", strerror(errno));

#ifdef SO_REUSEPORT
    if (setsockopt(listen_sock, SOL_SOCKET, SO_REUSEPORT, (int[1]){ 1 }, sizeof(int)) < 0)
        warn("setsockopt(SO_REUSEPORT): %s", strerror(errno));
#endif

    if (bind(listen_sock, (struct sockaddr *)&sa, salen) < 0)
        croak("bind: %s", strerror(errno));

    if (listen(listen_sock, SOMAXCONN) < 0)
        croak("listen: %s", strerror(errno));

    return listen_sock;
}

void setup_sighandler(int restartable)
{
    struct sigaction sa;

    sa.sa_handler = handle_hup;
    sigemptyset(&sa.sa_mask);
    sa.sa_flags = restartable ? SA_RESTART : 0;

    sigaction(SIGHUP, &sa, NULL);
}

void call_all(char *func)
{
    CV *cb;
    int i;

    for (i = 0; i < namespace_count; i++)
        if ((cb = get_cv(form("%s::%s", namespaces[i], func), FALSE)))
            callback_noargs(cb);
}

SV *new_state(int sock)
{
    struct milter_ctx *ctx;
    SV *state;
    SV *rstate;

    if (!(ctx = malloc(sizeof(*ctx))))
        croak("malloc");

    bzero(ctx, sizeof(*ctx));
    ctx->macro = newHV();
    ctx->sock = sock;

    state = newSViv(PTR2IV(ctx));
    rstate = sv_bless(newRV_inc(state), gv_stashpv("Milter", TRUE));
    SvREFCNT_dec(state);

    return rstate;
}

void free_state(SV *state)
{
    struct milter_ctx *ctx = MILTER_CTX(state);

    if (ctx->priv && SvREFCNT(ctx->priv) != 1)
        warn("REFCNT: ctx->priv=%d\n", SvREFCNT(ctx->priv));
    if (SvREFCNT(ctx->macro) != 1)
        warn("REFCNT: ctx->macro=%d\n", SvREFCNT(ctx->macro));
    if (SvREFCNT(state) != 1)
        warn("REFCNT: state=%d\n", SvREFCNT(state));

    if (ctx->priv)
        SvREFCNT_dec(ctx->priv);
    SvREFCNT_dec(ctx->macro);
    SvREFCNT_dec(state);
    free(ctx);
}

void process_child(int listen_sock)
{
    int req = 100 + (rand() % 250);

    call_all("init");

    while (req-- > 0 && !hup)
    {
        struct sockaddr sa;
        socklen_t salen = sizeof(sa);
        int sock;
        SV *state;

        setup_sighandler(0);
        if ((sock = accept(listen_sock, &sa, &salen)) < 0)
            continue;
        setup_sighandler(1);

        state = new_state(sock);
        for (;;)
            if (process_packet(state) < 0)
                break;
        free_state(state);

        close(sock);
    }

    call_all("cleanup");

    perl_destruct(my_perl);
    perl_free(my_perl);

    exit(0);
}

void pre_fork(int listen_sock, int maxchildren)
{
    pid_t pid;
    int i;

    for (i = 0; i < maxchildren; i++)
    {
        if ((pid = perl_fork()) < 0)
            croak("fork: %s", strerror(errno));

        if (!pid) /* child */
            process_child(listen_sock);
    }

    while (!hup)
    {
        int status = 0;

        setup_sighandler(0);
        if (wait(&status) < 0)
            continue;
        setup_sighandler(1);

        if ((pid = perl_fork()) < 0)
            croak("fork: %s", strerror(errno));

        if (!pid) /* child */
            process_child(listen_sock);
    }
}

void accept_and_fork(int listen_sock)
{
    struct sigaction sa;

    sa.sa_handler = SIG_IGN;
    sigemptyset(&sa.sa_mask);
    sigaction(SIGCHLD, &sa, NULL);

    while (!hup)
    {
        struct sockaddr sa;
        socklen_t salen = sizeof(sa);
        int sock;
        pid_t pid;
        SV *state;

        if ((sock = accept(listen_sock, &sa, &salen)) < 0)
            continue;

        if ((pid = perl_fork()) < 0)
            croak("fork: %s", strerror(errno));

        if (pid)
        {
            close(sock);
            continue;
        }

        call_all("init");

        state = new_state(sock);
        for (;;)
            if (process_packet(state) < 0)
                break;
        free_state(state);

        call_all("cleanup");

        perl_destruct(my_perl);
        perl_free(my_perl);

        exit(0);
    }
}

void usage(char *cmd)
{
    croak("Usage: %s [options] script...\n"
    "  -d         debugging output (+foreground)\n"
    "  -f         foreground\n"
    "  -l str     inet:<port>@<address> or unix:<path>\n"
    "  -n threads 0=accept and fork, >0=prefork N\n"
    "  -u uid     setuid to account\n"
    "  -g gid     setgid to group\n",
        cmd);
}

int main(int argc, char *argv[])
{
    char *cmd;
    char *embedding[] = { NULL, "-w", "-e", "0" };
    char *listen_str = DEFAULT_LISTEN;
    int threads = DEFAULT_THREADS;
    uid_t uid = 0;
    gid_t gid = 0;
    int ch;
    int listen_sock;
    int foreground = 0;

    cmd = strchr(argv[0], '/');
    cmd = cmd ? cmd + 1 : argv[0];
    embedding[0] = cmd;

    my_perl = perl_alloc();
    perl_construct(my_perl);
    perl_parse(my_perl, xs_init, sizeof(embedding)/sizeof(*embedding), embedding, NULL);

    while ((ch = getopt(argc, argv, "dfl:n:u:g:")) != -1)
    {
        switch (ch) {
            case 'd':
                debug = 1;
            case 'f': /* fallthrough */
                foreground = 1;
                break;
            case 'l':
                listen_str = optarg;
                break;
            case 'n':
                threads = atoi(optarg);
                break;
            case 'u': {
                char *foo = optarg;
                struct passwd *pwent;
                if ((pwent = getpwnam(foo)))
                    uid = pwent->pw_uid;
                else if (!(uid = atoi(foo)))
                    croak("unknown user/uid: %s\n", foo);
            } break;
            case 'g': {
                char *foo = optarg;
                struct group *grent;
                if ((grent = getgrnam(foo)))
                    gid = grent->gr_gid;
                else if (!(gid = atoi(foo)))
                    croak("unknown group/gid: %s\n", foo);
            } break;
            case '?':
            default:
                usage(cmd);
        }
    }

    if (optind >= argc)
        usage(cmd);

    PL_exit_flags |= PERL_EXIT_DESTRUCT_END;

    for (; optind < argc; optind++)
    {
        char namespace[strlen(argv[optind]) + 1];
        char line[2048];
        char *ptr = namespace;
        FILE *fp;
        SV *sv;
        int i;

        for (i = (argv[optind][0] == '/' ? 1 : 0); argv[optind][i]; i++)
            *ptr++ = isALNUM(argv[optind][i]) ? argv[optind][i] : '_';
        *ptr = 0;

        namespaces = realloc(namespaces, (namespace_count + 1) * sizeof(*namespaces));
        namespaces[namespace_count++] = strdup(namespace);

#define EXPORT_CONSTANT(name) do {                  \
    HV *stash = gv_stashpv(namespace, TRUE);        \
    (void)newCONSTSUB(stash, #name, newSViv(name)); \
} while (0)

        EXPORT_CONSTANT(SMFIS_CONTINUE);
        EXPORT_CONSTANT(SMFIS_REJECT);
        EXPORT_CONSTANT(SMFIS_DISCARD);
        EXPORT_CONSTANT(SMFIS_ACCEPT);
        EXPORT_CONSTANT(SMFIS_TEMPFAIL);

        EXPORT_CONSTANT(SMFIA_UNKNOWN);
        EXPORT_CONSTANT(SMFIA_UNIX);
        EXPORT_CONSTANT(SMFIA_INET);
        EXPORT_CONSTANT(SMFIA_INET6);

        sv = newSVpvf("package %s;\n# line 1 %s\n", namespace, argv[optind]);

        if (!(fp = fopen(argv[optind], "r")))
            croak("couldn't open %s", argv[optind]);
        while ((ptr = fgets(line, sizeof(line), fp)))
            sv_catpvn(sv, ptr, strlen(ptr));
        fclose(fp);

        Trace("Loading %s into %s\n", argv[optind], namespace);
            (void)eval_sv(sv, G_VOID|G_NOARGS|G_DISCARD);

        callback_init(namespace);

        SvREFCNT_dec(sv);
    }

    umask(0077);

    if (gid != 0)
    {
        if (setgroups(1, (gid_t *)&(gid_t[2]){ gid, 0 }))
            croak("setgroups: %s", strerror(errno));

        if (setgid(gid))
            croak("setgid: %s", strerror(errno));

        if (setegid(gid))
            croak("setegid: %s", strerror(errno));
    }

    if (uid != 0)
    {
        if (setuid(uid))
            croak("setuid: %s", strerror(errno));

        if (seteuid(uid))
            croak("seteuid: %s", strerror(errno));
    }

    if (!strncmp(listen_str, "inet:", sizeof("inet:")-1))
    {
        struct in_addr addr;
        uint16_t port;

        parse_address(listen_str + sizeof("inet:") - 1, &addr, &port);

        listen_sock = bind_server_inet(&addr, port);
    }
    else if (!strncmp(listen_str, "unix:", sizeof("unix:")-1))
    {
        char *path = listen_str + sizeof("unix:") - 1;

        if (1)
        {
            static int lock_fd = -1;
            char lock_file[PATH_MAX] = { 0 }, *p;
            strcpy(lock_file, path);
            if (!(p = rindex(lock_file, '/')))
                croak("unable to find '/' in path\n");
            *p = 0;
            strcat(lock_file, "/lock");
            if ((lock_fd = open(lock_file, O_CREAT|O_RDWR, 0600)) < 0)
                croak("Unable to open lock file: %s\n", lock_file);
            if (flock(lock_fd, LOCK_EX|LOCK_NB) < 0)
                croak("Unable to lock exclusively\n");
        }

        listen_sock = bind_server_unix(path);
    }
    else
        croak("invalid listen_str: %s", listen_str);

    if (!foreground && daemon(0, 0) < 0)
        warn("daemon: %s", strerror(errno));

    if (threads)
        pre_fork(listen_sock, threads);
    else
        accept_and_fork(listen_sock);

    perl_destruct(my_perl);
    perl_free(my_perl);

    return 0;
}
